package health

import (
	"context"
	"google.golang.org/grpc/health/grpc_health_v1"
	"sync/atomic"
)

type (
	GrpcHealthCheckFuc func(ctx context.Context, service string) (*grpc_health_v1.HealthCheckResponse, error)

	// GrpcHealthServer rpc 健康检查
	GrpcHealthServer interface {
		grpc_health_v1.HealthServer
		// Shutdown 服务关闭
		Shutdown()
		// Resume 服务恢复
		Resume()
	}

	// GrpcServer implements `GrpcHealthServer`.
	GrpcServer struct {
		grpc_health_v1.UnimplementedHealthServer
		// If shutdown 为 true，则所有服务状态均为 NOT_SERVING，并将保持 NOT_SERVING
		shutdown atomic.Bool
		checkFuc GrpcHealthCheckFuc
	}
)

func NewServer(checkFuc GrpcHealthCheckFuc) *GrpcServer {
	s := &GrpcServer{
		checkFuc: checkFuc,
	}
	s.shutdown.Store(false)
	return s
}

// Check implements `service Health`.
func (s *GrpcServer) Check(ctx context.Context, in *grpc_health_v1.HealthCheckRequest) (*grpc_health_v1.HealthCheckResponse, error) {
	if s.shutdown.Load() {
		/* 服务已关闭 */
		return &grpc_health_v1.HealthCheckResponse{
			Status: grpc_health_v1.HealthCheckResponse_NOT_SERVING,
		}, nil
	}
	if s.checkFuc == nil {
		return &grpc_health_v1.HealthCheckResponse{
			Status: grpc_health_v1.HealthCheckResponse_SERVING,
		}, nil
	}
	return s.checkFuc(ctx, in.Service)
}

func (s *GrpcServer) Shutdown() {
	s.shutdown.Store(true)
}

func (s *GrpcServer) Resume() {
	s.shutdown.Store(false)
}
